from abc import ABCMeta, abstractmethod
import networkx as nx


class GraphBase(metaclass=ABCMeta):
    """ 图构建基类 """

    def __init__(self, *args):
        super(GraphBase, self).__init__(*args)

    def prepare_graph(self, ids, batch_size=100, is_direct=False, is_weight=False):
        ''' 构建图
        :param ids: list
        :param is_direct:
        :param is_weight:
        :return:
        '''

        def chunks(l, batch):
            for i in range(0, len(l), batch):
                yield l[i:i + batch]

        G = nx.Graph()
        if is_direct:
            G = nx.DiGraph()

        for i, each in enumerate(chunks(ids, batch_size)):
            edge_list = self.create_edge(each)
            if is_weight:
                G.add_weighted_edges_from(edge_list)
            else:
                G.add_edges_from(edge_list)
        return G

    @abstractmethod
    def create_edge(self, batchids):
        '''
        :param batchids:
        :return: tuple of list, such as [(1, 2, 0.5),...] which means (start_node,end_node,weight)
        '''
        # read data from json or neo4j
        pass
